{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "anchors are:\n",
      "5.699999999999999,6.6499999999999995, 10.924999999999999,9.025, 8.075,16.15, 15.2,14.725, 23.275,20.425, 37.525,27.549999999999997, 50.824999999999996,43.225, 91.19999999999999,64.125, 171.0,145.35\n",
      "the average iou is:\n",
      "0.7033912071977335\n"
     ]
    }
   ],
   "source": [
    "\n",
    "from __future__ import division, print_function\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "def iou(box, clusters):\n",
    "    \"\"\"\n",
    "    Calculates the Intersection over Union (IoU) between a box and k clusters.\n",
    "    param:\n",
    "        box: tuple or array, shifted to the origin (i. e. width and height)\n",
    "        clusters: numpy array of shape (k, 2) where k is the number of clusters\n",
    "    return:\n",
    "        numpy array of shape (k, 0) where k is the number of clusters\n",
    "    \"\"\"\n",
    "    x = np.minimum(clusters[:, 0], box[0])\n",
    "    y = np.minimum(clusters[:, 1], box[1])\n",
    "    if np.count_nonzero(x == 0) > 0 or np.count_nonzero(y == 0) > 0:\n",
    "        raise ValueError(\"Box has no area\")\n",
    "\n",
    "    intersection = x * y\n",
    "    box_area = box[0] * box[1]\n",
    "    cluster_area = clusters[:, 0] * clusters[:, 1]\n",
    "\n",
    "    iou_ = np.true_divide(intersection, box_area + cluster_area - intersection + 1e-10)\n",
    "    # iou_ = intersection / (box_area + cluster_area - intersection + 1e-10)\n",
    "\n",
    "    return iou_\n",
    "\n",
    "\n",
    "def avg_iou(boxes, clusters):\n",
    "    \"\"\"\n",
    "    Calculates the average Intersection over Union (IoU) between a numpy array of boxes and k clusters.\n",
    "    param:\n",
    "        boxes: numpy array of shape (r, 2), where r is the number of rows\n",
    "        clusters: numpy array of shape (k, 2) where k is the number of clusters\n",
    "    return:\n",
    "        average IoU as a single float\n",
    "    \"\"\"\n",
    "    return np.mean([np.max(iou(boxes[i], clusters)) for i in range(boxes.shape[0])])\n",
    "\n",
    "\n",
    "def translate_boxes(boxes):\n",
    "    \"\"\"\n",
    "    Translates all the boxes to the origin.\n",
    "    param:\n",
    "        boxes: numpy array of shape (r, 4)\n",
    "    return:\n",
    "    numpy array of shape (r, 2)\n",
    "    \"\"\"\n",
    "    new_boxes = boxes.copy()\n",
    "    for row in range(new_boxes.shape[0]):\n",
    "        new_boxes[row][2] = np.abs(new_boxes[row][2] - new_boxes[row][0])\n",
    "        new_boxes[row][3] = np.abs(new_boxes[row][3] - new_boxes[row][1])\n",
    "    return np.delete(new_boxes, [0, 1], axis=1)\n",
    "\n",
    "\n",
    "def kmeans(boxes, k, dist=np.median):\n",
    "    \"\"\"\n",
    "    Calculates k-means clustering with the Intersection over Union (IoU) metric.\n",
    "    param:\n",
    "        boxes: numpy array of shape (r, 2), where r is the number of rows\n",
    "        k: number of clusters\n",
    "        dist: distance function\n",
    "    return:\n",
    "        numpy array of shape (k, 2)\n",
    "    \"\"\"\n",
    "    rows = boxes.shape[0]\n",
    "\n",
    "    distances = np.empty((rows, k))\n",
    "    last_clusters = np.zeros((rows,))\n",
    "\n",
    "    np.random.seed()\n",
    "\n",
    "    # the Forgy method will fail if the whole array contains the same rows\n",
    "    clusters = boxes[np.random.choice(rows, k, replace=False)]\n",
    "\n",
    "    while True:\n",
    "        for row in range(rows):\n",
    "            distances[row] = 1 - iou(boxes[row], clusters)\n",
    "\n",
    "        nearest_clusters = np.argmin(distances, axis=1)\n",
    "\n",
    "        if (last_clusters == nearest_clusters).all():\n",
    "            break\n",
    "\n",
    "        for cluster in range(k):\n",
    "            clusters[cluster] = dist(boxes[nearest_clusters == cluster], axis=0)\n",
    "\n",
    "        last_clusters = nearest_clusters\n",
    "\n",
    "    return clusters\n",
    "\n",
    "\n",
    "def parse_anno(annotation_path, target_size=None):\n",
    "    anno = open(annotation_path, 'r')\n",
    "    result = []\n",
    "    img_w = 1280\n",
    "    img_h = 720\n",
    "    i = 0\n",
    "    for line in anno:\n",
    "        i +=1\n",
    "        if i == 2000:\n",
    "            break\n",
    "        s = line.strip().split(' ')\n",
    "        s = s[2:]\n",
    "        box_cnt = len(s) // 5\n",
    "        for i in range(box_cnt):\n",
    "            tem = s[i].split(',')\n",
    "            x_min, y_min, x_max, y_max = float(tem[0]), float(tem[1]), float(tem[2]), float(tem[3])\n",
    "            width = x_max - x_min\n",
    "            height = y_max - y_min\n",
    "            assert width > 0\n",
    "            assert height > 0\n",
    "            # use letterbox resize, i.e. keep the original aspect ratio\n",
    "            # get k-means anchors on the resized target image size\n",
    "            if target_size is not None:\n",
    "                resize_ratio = min(target_size[0] / img_w, target_size[1] / img_h)\n",
    "                width *= resize_ratio\n",
    "                height *= resize_ratio\n",
    "                result.append([width, height])\n",
    "            # get k-means anchors on the original image size\n",
    "            else:\n",
    "                result.append([width, height])\n",
    "    result = np.asarray(result)\n",
    "    return result\n",
    "\n",
    "\n",
    "def get_kmeans(anno, cluster_num=9):\n",
    "\n",
    "    anchors = kmeans(anno, cluster_num)\n",
    "    ave_iou = avg_iou(anno, anchors)\n",
    "\n",
    "    anchors = anchors.astype('float').tolist()\n",
    "\n",
    "    anchors = sorted(anchors, key=lambda x: x[0] * x[1])\n",
    "\n",
    "    return anchors, ave_iou\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    # target resize format: [width, height]\n",
    "    # if target_resize is speficied, the anchors are on the resized image scale\n",
    "    # if target_resize is set to None, the anchors are on the original image scale\n",
    "    target_size = [608, 608]\n",
    "    annotation_path = \"C:/Users/ZZK/Documents/auto_drive/train_dataset/train.txt\"\n",
    "    anno_result = parse_anno(annotation_path, target_size=target_size)\n",
    "    anchors, ave_iou = get_kmeans(anno_result, 9)\n",
    "\n",
    "    anchor_string = ''\n",
    "    for anchor in anchors:\n",
    "        anchor_string += '{},{}, '.format(anchor[0], anchor[1])\n",
    "    anchor_string = anchor_string[:-2]\n",
    "\n",
    "    print('anchors are:')\n",
    "    print(anchor_string)\n",
    "    print('the average iou is:')\n",
    "    print(ave_iou)\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
